class Solution {
    public long countPairs(int n, int[][] edges) {
        long pairs = 0;
        UnionFind uf = new UnionFind(n);
        for (int[] edge : edges) {
            uf.union(edge[0], edge[1]);
        }
        long prevCount = 0;
        Map<Integer, Integer> sizes = new HashMap<Integer, Integer>();
        for (int i = 0; i < n; i++) {
            int root = uf.find(i);
            sizes.put(root, sizes.getOrDefault(root, 0) + 1);
        }
        Set<Map.Entry<Integer, Integer>> entries = sizes.entrySet();
        for (Map.Entry<Integer, Integer> entry : entries) {
            long currCount = entry.getValue();
            pairs += prevCount * currCount;
            prevCount += currCount;
        }
        return pairs;
    }
}

class UnionFind {
    private int[] parent;
    private int[] rank;

    public UnionFind(int n) {
        parent = new int[n];
        for (int i = 0; i < n; i++) {
            parent[i] = i;
        }
        rank = new int[n];
    }

    public void union(int x, int y) {
        int rootx = find(x);
        int rooty = find(y);
        if (rootx != rooty) {
            if (rank[rootx] > rank[rooty]) {
                parent[rooty] = rootx;
            } else if (rank[rootx] < rank[rooty]) {
                parent[rootx] = rooty;
            } else {
                parent[rooty] = rootx;
                rank[rootx]++;
            }
        }
    }

    public int find(int x) {
        if (parent[x] != x) {
            parent[x] = find(parent[x]);
        }
        return parent[x];
    }
}
